package com.whz.generic.crypt.rsa;

import lombok.SneakyThrows;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * 秘钥对长度为2048bit的RSA加密解密
 * <p>
 * 一般来讲公钥加密,私钥解密;反过来也可以
 * </p>
 *
 * @author whz
 */
public class RsaUtils {
    // 秘钥最大长度 256
    private static final int KEY_SIZE = 2 * 1024;
    // 解密的密文最大字节
    private static final int MAX_DEC_LENGTH = (KEY_SIZE >> 3);
    // 加密的明文最大字节,还要减去额外的签名长度
    private static final int MAX_ENC_LENGTH = MAX_DEC_LENGTH - 11;
    // 算法名称
    private static final String RSA = "RSA";

    /**
     * 随机生成RSA的密钥对,秘钥已经base64编码
     */
    @SneakyThrows
    public static RsaKeyPair getKeyPair() {
        //KeyPairGenerator类用于生成公钥和密钥对，基于RSA算法生成对象
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(RSA);
        //初始化密钥对生成器，密钥大小为96-1024位
        keyPairGen.initialize(KEY_SIZE, new SecureRandom());
        //生成一个密钥对，保存在keyPair中
        KeyPair keyPair = keyPairGen.generateKeyPair();
        //得到私钥
        PrivateKey privateKey = keyPair.getPrivate();
        //得到公钥
        PublicKey publicKey = keyPair.getPublic();
        //得到公钥字符串
        String publicKeyString = Base64.getEncoder().encodeToString(publicKey.getEncoded());
        //得到私钥字符串
        String privateKeyString = Base64.getEncoder().encodeToString(privateKey.getEncoded());
        //将公钥和私钥保存
        return new RsaKeyPair(publicKeyString, privateKeyString);
    }

    /**
     * RSA公钥加密
     *
     * @param originMessage 加密字符串,UTF-8 编码
     * @param publicKey     base64编码公钥
     * @return 密文
     */
    @SneakyThrows
    public static String encryptByPublicKey(String originMessage, String publicKey) {
        RSAPublicKey pubKey = getRsaPublicKeyFromString(publicKey);
        return encrypt(originMessage.getBytes(StandardCharsets.UTF_8), pubKey);
    }


    /**
     * RSA私钥加密
     *
     * @param originMessage 加密字符串,UTF-8 编码
     * @param privateKey    base64编码的公钥
     * @return 密文
     */
    @SneakyThrows
    public static String encryptByPrivateKey(String originMessage, String privateKey) {
        PrivateKey priKey = getPrivateKeyFromString(privateKey);
        return encrypt(originMessage.getBytes(StandardCharsets.UTF_8), priKey);
    }


    /**
     * 公钥加密
     *
     * @param originMessage 原始文本
     * @param pubKey        公钥
     * @return 密文BASE64编码
     */
    @SneakyThrows
    public static String encrypt(byte[] originMessage, RSAPublicKey pubKey) {
        //RAS加密
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.ENCRYPT_MODE, pubKey);
        return Base64.getEncoder().encodeToString(doFinal(cipher, originMessage, MAX_ENC_LENGTH));
    }

    /**
     * 私钥加密
     *
     * @param originMessage 原始文本
     * @param privateKey    私钥
     * @return 密文BASE64编码
     */
    @SneakyThrows
    public static String encrypt(byte[] originMessage, PrivateKey privateKey) {
        //RAS加密
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.ENCRYPT_MODE, privateKey);
        return Base64.getEncoder().encodeToString(doFinal(cipher, originMessage, MAX_ENC_LENGTH));
    }

    /**
     * RSA私钥解密
     *
     * @param decryptMessage 加密字符串,BASE64 编码
     * @param privateKey     base64编码的私钥
     * @return 明文
     */
    @SneakyThrows
    public static String decryptByPrivateKey(String decryptMessage, String privateKey) {
        PrivateKey priKey = getPrivateKeyFromString(privateKey);
        return decrypt(Base64.getDecoder().decode(decryptMessage), priKey);
    }

    /**
     * RSA公钥解密
     *
     * @param decryptMessage 加密字符串,BASE64 编码
     * @param publicKey      base64编码的公钥
     * @return 明文
     */
    @SneakyThrows
    public static String decryptByPublicKey(String decryptMessage, String publicKey) {
        RSAPublicKey pubKey = getRsaPublicKeyFromString(publicKey);
        return decrypt(Base64.getDecoder().decode(decryptMessage), pubKey);
    }

    /**
     * 公钥解密
     *
     * @param decryptMessage 密文
     * @param publicKey      公钥
     * @return 明文UTF-8编码
     */
    @SneakyThrows
    public static String decrypt(byte[] decryptMessage, RSAPublicKey publicKey) {
        //RSA解密
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.DECRYPT_MODE, publicKey);
        return new String(doFinal(cipher, decryptMessage, MAX_DEC_LENGTH), StandardCharsets.UTF_8);
    }

    /**
     * 私钥解密
     *
     * @param decryptMessage 密文
     * @param privateKey     私钥
     * @return 明文UTF-8编码
     */
    @SneakyThrows
    public static String decrypt(byte[] decryptMessage, PrivateKey privateKey) {
        //RSA解密
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return new String(doFinal(cipher, decryptMessage, MAX_DEC_LENGTH), StandardCharsets.UTF_8);
    }


    /**
     * base64的公钥转换
     *
     * @param publicKey 公钥的base64编码字符串
     * @return X509编码的公钥
     */
    @SneakyThrows
    private static RSAPublicKey getRsaPublicKeyFromString(String publicKey) {
        //base64编码的公钥
        byte[] decoded = Base64.getDecoder().decode(publicKey);
        return (RSAPublicKey) KeyFactory.getInstance(RSA).generatePublic(new X509EncodedKeySpec(decoded));
    }

    /**
     * base64编码的私钥转换
     *
     * @param privateKey base64编码的私钥
     * @return PKCS8编码的私钥
     */
    @SneakyThrows
    public static PrivateKey getPrivateKeyFromString(String privateKey) {
        //Base64编码的私钥
        byte[] decoded = Base64.getDecoder().decode(privateKey);
        return KeyFactory.getInstance(RSA).generatePrivate(new PKCS8EncodedKeySpec(decoded));
    }

    /**
     * 支持大文本加密和解密,分片加密
     */
    @SneakyThrows
    private static byte[] doFinal(Cipher cipher, byte[] input, int maxEncLength) {
        int inputLen = input.length;
        try (ByteArrayOutputStream out = new ByteArrayOutputStream();) {
            int offSet = 0;
            byte[] cache;
            // Sectional Encryption of Data
            while (inputLen - offSet > 0) {
                if (inputLen - offSet >= maxEncLength) {
                    cache = cipher.doFinal(input, offSet, maxEncLength);
                } else {
                    cache = cipher.doFinal(input, offSet, inputLen - offSet);
                }
                out.write(cache, 0, cache.length);
                offSet += maxEncLength;
            }
            return out.toByteArray();
        }
    }
}